// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/base/pollux_exception.h>
#include <pollux/testing/gtest_utils.h>
#include <pollux/expression/vector_function.h>
#include <pollux/functions/macros.h>
#include <pollux/functions/registerer.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>
#include <pollux/type/opaque_custom_types.h>

namespace kumo::pollux::test {

class CustomTypeTest : public functions::test::FunctionBaseTest {};

namespace {
struct FancyInt {
  const int64_t n;

  explicit FancyInt(int64_t _n) : n{_n} {}
};

class FancyIntType : public OpaqueType {
  FancyIntType() : OpaqueType(std::type_index(typeid(FancyInt))) {}

 public:
  static const std::shared_ptr<const FancyIntType>& get() {
    static const std::shared_ptr<const FancyIntType> instance{
        new FancyIntType()};

    return instance;
  }

  std::string toString() const override {
    return name();
  }

  const char* name() const override {
    return "fancy_int";
  }
};

class FancyIntTypeFactories : public CustomTypeFactories {
 public:
  TypePtr getType(const std::vector<TypeParameter>& parameters) const override {
    POLLUX_CHECK(parameters.empty());
    return FancyIntType::get();
  }

  exec::CastOperatorPtr getCastOperator() const override {
    POLLUX_UNSUPPORTED();
  }

  AbstractInputGeneratorPtr getInputGenerator(
      const InputGeneratorConfig& /*config*/) const override {
    return nullptr;
  }
};

class ToFancyIntFunction : public exec::VectorFunction {
 public:
  void apply(
      const SelectivityVector& rows,
      std::vector<VectorPtr>& args,
      const TypePtr& outputType,
      exec::EvalCtx& context,
      VectorPtr& result) const override {
    auto flatInput = args[0]->as<SimpleVector<int64_t>>();

    BaseVector::ensure_writable(rows, outputType, context.pool(), result);
    auto flatResult = result->as_flat_vector<std::shared_ptr<void>>();

    rows.applyToSelected([&](auto row) {
      flatResult->set(row, std::make_shared<FancyInt>(flatInput->value_at(row)));
    });
  }

  static std::vector<std::shared_ptr<exec::FunctionSignature>> signatures() {
    // bigint -> fancy_int
    return {exec::FunctionSignatureBuilder()
                .returnType("fancy_int")
                .argumentType("bigint")
                .build()};
  }
};

class FromFancyIntFunction : public exec::VectorFunction {
 public:
  void apply(
      const SelectivityVector& rows,
      std::vector<VectorPtr>& args,
      const TypePtr& /* outputType */,
      exec::EvalCtx& context,
      VectorPtr& result) const override {
    auto flatInput = args[0]->as<SimpleVector<std::shared_ptr<void>>>();

    BaseVector::ensure_writable(rows, BIGINT(), context.pool(), result);
    auto flatResult = result->as_flat_vector<int64_t>();

    rows.applyToSelected([&](auto row) {
      flatResult->set(
          row, std::static_pointer_cast<FancyInt>(flatInput->value_at(row))->n);
    });
  }

  static std::vector<std::shared_ptr<exec::FunctionSignature>> signatures() {
    // fancy_int -> bigint
    return {exec::FunctionSignatureBuilder()
                .returnType("bigint")
                .argumentType("fancy_int")
                .build()};
  }
};

// Define type to use in simple functions.
struct FancyIntT {
  using type = std::shared_ptr<FancyInt>;
  static constexpr const char* typeName = "fancy_int";
};
using TheFancyInt = CustomType<FancyIntT>;

template <typename T>
struct FancyPlusFunction {
  POLLUX_DEFINE_FUNCTION_TYPES(T);

  void call(
      out_type<TheFancyInt>& result,
      const arg_type<TheFancyInt>& a,
      const arg_type<TheFancyInt>& b) {
    result = std::make_shared<FancyInt>(a->n + b->n);
  }
};

class AlwaysFailingTypeFactories : public CustomTypeFactories {
 public:
  TypePtr getType(const std::vector<TypeParameter>& parameters) const override {
    POLLUX_CHECK(parameters.empty());
    POLLUX_UNSUPPORTED();
  }

  exec::CastOperatorPtr getCastOperator() const override {
    POLLUX_UNSUPPORTED();
  }

  AbstractInputGeneratorPtr getInputGenerator(
      const InputGeneratorConfig& /*config*/) const override {
    POLLUX_UNSUPPORTED();
  }
};
} // namespace

/// Register custom type based on OpaqueType. Register a vector function that
/// produces this type, another vector function that consumes this type, and a
/// simple function that takes and returns this type. Verify function signatures
/// and evaluate some expressions.
TEST_F(CustomTypeTest, customType) {
  ASSERT_TRUE(registerCustomType(
      "fancy_int", std::make_unique<FancyIntTypeFactories>()));

  ASSERT_FALSE(registerCustomType(
      "fancy_int", std::make_unique<AlwaysFailingTypeFactories>()));

  register_function<FancyPlusFunction, TheFancyInt, TheFancyInt, TheFancyInt>(
      {"fancy_plus"});

  exec::registerVectorFunction(
      "to_fancy_int",
      ToFancyIntFunction::signatures(),
      std::make_unique<ToFancyIntFunction>());
  exec::registerVectorFunction(
      "from_fancy_int",
      FromFancyIntFunction::signatures(),
      std::make_unique<FromFancyIntFunction>());

  // Verify signatures.
  auto signatures = getSignatureStrings("fancy_plus");
  ASSERT_EQ(1, signatures.size());
  ASSERT_EQ(1, signatures.count("(fancy_int,fancy_int) -> fancy_int"));

  signatures = getSignatureStrings("to_fancy_int");
  ASSERT_EQ(1, signatures.size());
  ASSERT_EQ(1, signatures.count("(bigint) -> fancy_int"));

  signatures = getSignatureStrings("from_fancy_int");
  ASSERT_EQ(1, signatures.size());
  ASSERT_EQ(1, signatures.count("(fancy_int) -> bigint"));

  // Evaluate expressions.
  auto data = make_flat_vector<int64_t>({1, 2, 3, 4, 5});

  auto result =
      evaluate("from_fancy_int(to_fancy_int(c0))", make_row_vector({data}));
  assertEqualVectors(data, result);

  result = evaluate(
      "from_fancy_int(to_fancy_int(c0 + 10)) - 10", make_row_vector({data}));
  assertEqualVectors(data, result);

  result = evaluate(
      "from_fancy_int(fancy_plus(to_fancy_int(c0), to_fancy_int(10)))",
      make_row_vector({data}));
  auto expected = make_flat_vector<int64_t>({11, 12, 13, 14, 15});
  assertEqualVectors(expected, result);

  // Cleanup.
  ASSERT_TRUE(unregisterCustomType("fancy_int"));
  ASSERT_FALSE(unregisterCustomType("fancy_int"));
}

TEST_F(CustomTypeTest, getCustomTypeNames) {
  auto names = getCustomTypeNames();
  ASSERT_EQ(
      (std::unordered_set<std::string>{
          "JSON",
          "HYPERLOGLOG",
          "TIMESTAMP WITH TIME ZONE",
          "UUID",
          "IPADDRESS",
          "IPPREFIX",
          "BINGTILE",
          "TDIGEST",
          "GEOMETRY"}),
      names);

  ASSERT_TRUE(registerCustomType(
      "fancy_int", std::make_unique<FancyIntTypeFactories>()));

  names = getCustomTypeNames();
  ASSERT_EQ(
      (std::unordered_set<std::string>{
          "JSON",
          "HYPERLOGLOG",
          "TIMESTAMP WITH TIME ZONE",
          "UUID",
          "IPADDRESS",
          "IPPREFIX",
          "BINGTILE",
          "FANCY_INT",
          "TDIGEST",
          "GEOMETRY"}),
      names);

  ASSERT_TRUE(unregisterCustomType("fancy_int"));
}

TEST_F(CustomTypeTest, nullConstant) {
  ASSERT_TRUE(registerCustomType(
      "fancy_int", std::make_unique<FancyIntTypeFactories>()));
  auto checkNullConstant = [&](const TypePtr& type,
                               const std::string& expectedTypeString) {
    auto null = BaseVector::create_null_constant(type, 10, pool());
    EXPECT_TRUE(null->is_constant_encoding());
    EXPECT_TRUE(type->equivalent(*null->type()));
    EXPECT_EQ(type->toString(), null->type()->toString());
    EXPECT_EQ(type->toString(), expectedTypeString);
    for (auto i = 0; i < 10; ++i) {
      EXPECT_TRUE(null->is_null_at(i));
    }
  };
  auto names = getCustomTypeNames();
  for (const auto& name : names) {
    if (name == "TDIGEST") {
      auto type = getCustomType(name, {TypeParameter(DOUBLE())});
      checkNullConstant(type, "TDIGEST(DOUBLE)");
    } else {
      auto type = getCustomType(name, {});
      checkNullConstant(type, type->toString());
    }
  }
  ASSERT_TRUE(unregisterCustomType("fancy_int"));
}

struct Tuple {
  int64_t x;
  int64_t y;
};

static constexpr char kName[] = "tuple_type";
using TupleTypeRegistrar = OpaqueCustomTypeRegister<Tuple, kName>;
// The type used in the simple function interface.
using TupleType = typename TupleTypeRegistrar::SimpleType;

template <typename T>
struct FuncMake {
  POLLUX_DEFINE_FUNCTION_TYPES(T);

  void call(out_type<TupleType>& result, int64_t a, int64_t b) {
    result = std::make_shared<Tuple>(Tuple{a, b});
  }
};

template <typename T>
struct FuncPlus {
  POLLUX_DEFINE_FUNCTION_TYPES(T);
  void call(
      out_type<TupleType>& result,
      const arg_type<TupleType>& a,
      const arg_type<TupleType>& b) {
    result = std::make_shared<Tuple>(Tuple{a->x + b->x, a->y + b->y});
  }
};

template <typename T>
struct FuncReduce {
  POLLUX_DEFINE_FUNCTION_TYPES(T);
  void call(int64_t& result, const arg_type<TupleType>& a) {
    result = a->x + a->y;
  }
};

TEST_F(CustomTypeTest, testOpaqueCustomTypeAutoCreation) {
  // register untyped version.
  register_function<FuncMake, std::shared_ptr<Tuple>, int64_t, int64_t>(
      {"make_tuple_untyped"});

  ASSERT_ANY_THROW((
      register_function<FuncMake, TupleType, int64_t, int64_t>({"make_tuple"})));

  TupleTypeRegistrar::registerType();
  register_function<FuncMake, TupleType, int64_t, int64_t>({"make_tuple"});
  register_function<FuncPlus, TupleType, TupleType, TupleType>({"plus_tuple"});
  register_function<FuncReduce, int64_t, TupleType>({"reduce_tuple"});

  // Verify signatures.
  {
    auto signatures = getSignatureStrings("plus_tuple");
    ASSERT_EQ(1, signatures.size());
    ASSERT_EQ(1, signatures.count("(tuple_type,tuple_type) -> tuple_type"));
  }

  {
    auto signatures = getSignatureStrings("make_tuple_untyped");
    ASSERT_EQ(1, signatures.size());
    ASSERT_EQ(1, signatures.count("(bigint,bigint) -> opaque"));
  }

  {
    auto signatures = getSignatureStrings("make_tuple");
    ASSERT_EQ(1, signatures.size());
    ASSERT_EQ(1, signatures.count("(bigint,bigint) -> tuple_type"));
  }

  {
    auto signatures = getSignatureStrings("reduce_tuple");
    ASSERT_EQ(1, signatures.size());
    ASSERT_EQ(1, signatures.count("(tuple_type) -> bigint"));
  }

  auto data = make_flat_vector<int64_t>({1, 2, 3});

  // Evaluate expressions.
  {
    auto expected = make_flat_vector<int64_t>({4, 8, 12});
    auto result = evaluate(
        "reduce_tuple(plus_tuple(make_tuple(c0, c0), make_tuple(c0, c0)))",
        make_row_vector({data}));
    test::assertEqualVectors(expected, result);
  }

  {
    auto expected = make_flat_vector<int64_t>({2, 4, 6});
    auto result =
        evaluate("reduce_tuple(make_tuple(c0, c0))", make_row_vector({data}));
    test::assertEqualVectors(expected, result);
  }

  POLLUX_ASSERT_THROW(
      evaluate(
          "reduce_tuple(make_tuple_untyped(c0, c0))", make_row_vector({data})),
      "");
}
} // namespace kumo::pollux::test
